import scipy.io.wavfile
import numpy as np
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import seaborn as sns

BASE_DIR, FILENAME = os.path.split(__file__)
path = '../../../../../large_data/audio/zsn-stop1.wav'
AUDIO_PATH = os.path.join(BASE_DIR, path)

sep('Load')
rs, signal = scipy.io.wavfile.read(AUDIO_PATH)
print('rs', rs)
print_numpy_ndarray_info_simple(signal, 'signal')
if len(signal.shape) >= 2:
    signal = signal[:, 0]
print_numpy_ndarray_info_simple(signal, 'signal[:, 0]')

spr = 2
spc = 1
spn = 0
plt.figure(figsize=[6, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Signal')
plt.plot(signal)

sep('Emphasize')
signal = np.append(signal[0:1], signal[1:] - signal[:-1] * 0.97)
print_numpy_ndarray_info_simple(signal, 'signal emphasized')

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Signal emphasized')
plt.plot(signal)
plt.show()

sep('Frames')
signal_len = len(signal)
frame_len = int(np.ceil(rs / 40))
frame_stride = int(np.ceil(rs / 100))
n_frames = int(np.ceil((abs(signal_len - frame_len) + 1) / frame_stride))
padded_len = frame_len + frame_stride * n_frames
print('signal_len', signal_len)
print('frame_len', frame_len)
print('frame_stride', frame_stride)
print('padded_len', padded_len)
signal = np.append(signal, np.zeros(padded_len - signal_len, dtype=signal.dtype))
print_numpy_ndarray_info_simple(signal, 'signal padded')

idx_col = np.arange(frame_len).reshape(1, -1)
idx_row = np.arange(0, frame_stride * n_frames, frame_stride).reshape(-1, 1)
idx_frames = idx_col + idx_row
print(idx_frames)
print_numpy_ndarray_info_simple(idx_frames, 'idx_frames')
frames = signal[idx_frames]
print_numpy_ndarray_info_simple(frames, 'frames')

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('frames')
sns.heatmap(frames)

sep('Window')
win = np.hamming(frame_len)
frames *= win

spn += 1
plt.subplot(spr, spc, spn)
plt.title('frames windowed')
sns.heatmap(frames)
plt.show()

sep('Fourier trans')
NFFT = 512
rfft = np.fft.rfft(frames, n=NFFT)
print_numpy_ndarray_info_simple(rfft, 'rfft')
rfft = np.absolute(rfft)
print_numpy_ndarray_info_simple(rfft, 'rfft')
_, n_spec = rfft.shape
print('n_spec', n_spec)

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])
spn += 1
plt.subplot(spr, spc, spn)
plt.title('rfft')
sns.heatmap(rfft)

sep('Power spectrum')
power = rfft ** 2 / NFFT
print_numpy_ndarray_info_simple(power, 'power')

spn += 1
plt.subplot(spr, spc, spn)
plt.title('power')
sns.heatmap(power)
plt.show()
